import pickle
import tensorflow as tf
from sklearn.model_selection import train_test_split
import time
from sklearn.utils import shuffle

from alexnet import AlexNet


nb_classes = 43
epochs = 10
batch_size = 128

# TODO: Load traffic signs data.
with open('./train.p', 'rb') as f:
	data = pickle.load(f)

# TODO: Split data into training and validation sets.
X_train, X_val, y_train, y_val = train_test_split(data['features'], data['labels'], test_size=0.33, random_state = 0)

# TODO: Define placeholders and resize operation.
features = tf.placeholder(tf.float32, (None, 32, 32, 3))
labels = tf.placeholder(tf.int64, None)
resized = tf.image.resize_images(features, (227, 227))

# TODO: pass placeholder as first argument to `AlexNet`.
fc7 = AlexNet(resized, feature_extract=True)
# NOTE: `tf.stop_gradient` prevents the gradient from flowing backwards
# past this point, keeping the weights before and up to `fc7` frozen.
# This also makes training faster, less work to do!
# TODO: Add the final layer for traffic sign classification.
fc7 = tf.stop_gradient(fc7)
shape = (fc7.get_shape().as_list()[-1], nb_classes)
fc8W = tf.Variable(tf.truncated_normal(shape, stddev=1e-2))
fc8b = tf.Variable(tf.zeros(nb_classes))
logits = tf.nn.xw_plus_b(fc7, fc8W, fc8b)

# TODO: Define loss, training, accuracy operations.
# HINT: Look back at your traffic signs project solution, you may
# be able to reuse some the code.

cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=labels)
loss_op = tf.reduce_mean(cross_entropy)
opt = tf.train.AdamOptimizer()
train_op = opt.minimize(loss_op, var_list =[fc8W, fc8b])
init_op = tf.global_variables_initializer()

# TODO: Train and evaluate the feature extraction model.
preds = tf.argmax(logits, 1)
accuracy_op = tf.reduce_mean(tf.cast(tf.equal(preds, labels), tf.float32))

def eval_on_data(X, y, sess):
	total_acc = 0
	total_loss = 0
	for offset in range(0, X.shape[0], batch_size):
		end = offset + batch_size
		X_batch = X[offset:end]
		y_batch = y[offset:end]

		loss, acc = sess.run([loss_op, accuracy_op], feed_dict={features:X_batch, labels:y_batch})
		total_loss += (loss * X_batch.shape[0])
		total_acc += (acc * X_batch.shape[0])

	return total_loss/X.shape[0], total_acc/X.shape[0]

with tf.Session() as sess:
	sess.run(init_op)

	for i in range(epochs):
		# training
		X_train, y_train = shuffle(X_train, y_train)
		t0 = time.time()
		for offset in range(0, X_train.shape[0], batch_size):
			end = offset + batch_size
			sess.run(train_op, feed_dict={features: X_train[offset:end], labels:y_train[offset:end]})

		val_loss, val_acc = eval_on_data(X_val, y_val, sess)
		print("Epoch", i+1)
		print("Time: %.3f seconds" % (time.time()-t0))
		print("Validation Loss= ", val_loss)
		print("Validation Accuracy = ", val_acc)
		print("")


# (IntroToTensorFlow) C:\Users\alchemz\Documents\GitHub\Lab-AlexNet\CarND-Alexnet-Feature-Extraction>python train_feature_extraction.py
# 2018-02-23 12:49:53.091682: I C:\tf_jenkins\home\workspace\rel-win\M\windows\PY\36\tensorflow\core\platform\cpu_feature_guard.cc:137] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX AVX2
# Epoch 1
# Time: 1125.561 seconds
# Validation Loss=  0.505558593554
# Validation Accuracy =  0.867068552434

# Epoch 2
# Time: 2188.984 seconds
# Validation Loss=  0.389353002469
# Validation Accuracy =  0.886853698117

# Epoch 3
# Time: 1252.383 seconds
# Validation Loss=  0.264185377403
# Validation Accuracy =  0.935157276446

# Epoch 4
# Time: 1249.726 seconds
# Validation Loss=  0.216089001824
# Validation Accuracy =  0.945668135095

# Epoch 5
# Time: 1249.038 seconds
# Validation Loss=  0.184718909267
# Validation Accuracy =  0.955869850839

# Epoch 6
# Time: 1310.898 seconds
# Validation Loss=  0.173390239028
# Validation Accuracy =  0.955174279306

# Epoch 7
# Time: 1240.248 seconds
# Validation Loss=  0.152155665543
# Validation Accuracy =  0.962284566064

# Epoch 8
# Time: 1227.802 seconds
# Validation Loss=  0.14400710936
# Validation Accuracy =  0.961820851712

# Epoch 9
# Time: 1236.499 seconds
# Validation Loss=  0.136446398741
# Validation Accuracy =  0.963057423289

# Epoch 10
# Time: 1194.842 seconds
# Validation Loss=  0.122640631405
# Validation Accuracy =  0.967926423989